from __future__ import annotations

import os
from typing import Optional

import numpy as np
import pyqtgraph.opengl as gl
import trimesh
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtGui import QVector3D

from config.settings import settings


class PointCloudDialog(QtWidgets.QDialog):
    def __init__(self, parent=None, obj_data=None, groups=None) -> None:
        super().__init__(parent)
        self.setWindowTitle("点云采样")
        self.setGeometry(200, 200, 500, 400)
        self.obj_data = obj_data
        self.groups = groups
        self.point_cloud: Optional[dict] = None

        self.init_ui()

    def init_ui(self) -> None:
        layout = QtWidgets.QVBoxLayout(self)

        file_btn_layout = QtWidgets.QHBoxLayout()

        self.import_btn = QtWidgets.QPushButton("加载 OBJ")
        self.import_btn.setIcon(self.style().standardIcon(QtWidgets.QStyle.SP_DialogOpenButton))
        self.import_btn.clicked.connect(self.import_obj)
        file_btn_layout.addWidget(self.import_btn)

        self.sample_btn = QtWidgets.QPushButton("采样点云")
        self.sample_btn.clicked.connect(self.sample_point_cloud)
        file_btn_layout.addWidget(self.sample_btn)

        layout.addLayout(file_btn_layout)

        density_layout = QtWidgets.QHBoxLayout()
        density_layout.addWidget(QtWidgets.QLabel("点密度 (点/单位面积):"))
        self.density_spinbox = QtWidgets.QSpinBox()
        self.density_spinbox.setRange(1, 1000)
        self.density_spinbox.setValue(50)
        density_layout.addWidget(self.density_spinbox)
        layout.addLayout(density_layout)

        resolution_layout = QtWidgets.QHBoxLayout()
        resolution_layout.addWidget(QtWidgets.QLabel("采样分辨率:"))
        self.resolution_spinbox = QtWidgets.QSpinBox()
        self.resolution_spinbox.setRange(1, 1000)
        self.resolution_spinbox.setValue(100)
        resolution_layout.addWidget(self.resolution_spinbox)
        layout.addLayout(resolution_layout)

        self.rgb_checkbox = QtWidgets.QCheckBox("包含RGB颜色")
        self.rgb_checkbox.setChecked(True)
        layout.addWidget(self.rgb_checkbox)

        button_layout = QtWidgets.QHBoxLayout()

        self.visualize_btn = QtWidgets.QPushButton("可视化")
        self.visualize_btn.clicked.connect(self.visualize_point_cloud)
        self.visualize_btn.setEnabled(False)
        button_layout.addWidget(self.visualize_btn)

        self.save_btn = QtWidgets.QPushButton("保存点云")
        self.save_btn.clicked.connect(self.save_point_cloud)
        self.save_btn.setEnabled(False)
        button_layout.addWidget(self.save_btn)

        layout.addLayout(button_layout)

        self.status_label = QtWidgets.QLabel("就绪")
        self.status_label.setStyleSheet("color: #AAAAAA; font-size: 10px;")
        layout.addWidget(self.status_label)

        self.pc_view = gl.GLViewWidget()
        self.pc_view.setMinimumSize(400, 300)
        self.pc_view.setBackgroundColor(settings.viewer.background_color)
        layout.addWidget(self.pc_view)

    # ------------------------------------------------------------------ #
    # File operations
    # ------------------------------------------------------------------ #
    def import_obj(self) -> None:
        path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打开 OBJ 文件", "", "OBJ 文件 (*.obj)")
        if not path:
            return

        try:
            vertices = []
            faces = []

            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    if line.startswith("v "):
                        parts = line.strip().split()
                        vertices.append([float(p) for p in parts[1:]])
                    elif line.startswith("f "):
                        parts = line.strip().split()
                        face = []
                        for part in parts[1:]:
                            v_idx = int(part.split("/")[0]) - 1
                            face.append(v_idx)
                        faces.append(face)

            if not vertices or not faces:
                self.status_label.setText("无效的 OBJ 文件 - 缺少顶点或面数据")
                return

            self.obj_data = {"vertices": np.array(vertices), "faces": faces, "path": path}
            self.groups = None
            self.point_cloud = None
            self.pc_view.clear()

            self.status_label.setText(f"已导入: {os.path.basename(path)}")
            self.visualize_btn.setEnabled(False)
            self.save_btn.setEnabled(False)
        except Exception as exc:
            self.status_label.setText(f"导入失败: {exc}")

    # ------------------------------------------------------------------ #
    # Point cloud sampling
    # ------------------------------------------------------------------ #
    def sample_point_cloud(self) -> None:
        if not self.obj_data or not self.obj_data["faces"]:
            self.status_label.setText("错误: 没有可用的模型数据")
            return

        try:
            self.status_label.setText("正在采样点云...")
            QtWidgets.QApplication.processEvents()

            vertices = self.obj_data["vertices"]
            faces = self.obj_data["faces"]
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)

            point_density = self.density_spinbox.value()
            area = mesh.area
            num_points = int(area * point_density)

            num_points = max(100, min(num_points, 1_000_000))

            points, face_indices = trimesh.sample.sample_surface(mesh, num_points)

            colors = None
            if self.rgb_checkbox.isChecked():
                colors = np.zeros((len(points), 4))
                colors[:, :3] = [0.8, 0.8, 0.8]
                colors[:, 3] = 1.0

                if self.groups:
                    face_to_group = {}
                    for group_idx, group in enumerate(self.groups):
                        for face in group["faces"]:
                            for orig_face_idx, orig_face in enumerate(self.obj_data["faces"]):
                                if set(orig_face) == set(group["v_indices"][v] for v in face):
                                    face_to_group[orig_face_idx] = group_idx
                                    break

                    for i, face_idx in enumerate(face_indices):
                        if face_idx in face_to_group:
                            group_idx = face_to_group[face_idx]
                            colors[i] = self.groups[group_idx]["color"]

            self.point_cloud = {"points": points, "colors": colors, "num_points": num_points}

            self.status_label.setText(f"采样完成! 点云包含 {num_points} 个点")
            self.visualize_btn.setEnabled(True)
            self.save_btn.setEnabled(True)
        except Exception as exc:
            self.status_label.setText(f"采样失败: {exc}")

    def visualize_point_cloud(self) -> None:
        if not self.point_cloud:
            return

        self.pc_view.clear()

        points = self.point_cloud["points"]
        colors = self.point_cloud["colors"]

        point_cloud_item = gl.GLScatterPlotItem(
            pos=points,
            size=0.1,
            color=colors[:, :3] if colors is not None else (0.8, 0.8, 0.8, 1.0),
            pxMode=False,
        )

        self.pc_view.addItem(point_cloud_item)

        axis = gl.GLAxisItem()
        axis.setSize(10, 10, 10)
        self.pc_view.addItem(axis)

        if points.shape[0] > 0:
            min_coords = np.min(points, axis=0)
            max_coords = np.max(points, axis=0)
            size = np.max(max_coords - min_coords)
            center = (min_coords + max_coords) / 2
            distance = size * 2

            self.pc_view.setCameraPosition(
                pos=QVector3D(float(center[0]), float(center[1]), float(center[2] + distance)),
                distance=distance,
                elevation=30,
                azimuth=45,
            )

    def save_point_cloud(self) -> None:
        if not self.point_cloud:
            return

        default_path = ""
        if self.obj_data and "path" in self.obj_data:
            original_path = self.obj_data["path"]
            dirname = os.path.dirname(original_path)
            basename = os.path.splitext(os.path.basename(original_path))[0]
            default_path = os.path.join(dirname, f"{basename}_pointcloud.ply")

        path, _ = QtWidgets.QFileDialog.getSaveFileName(self, "保存点云", default_path, "点云文件 (*.ply)")
        if not path:
            return

        if not path.lower().endswith(".ply"):
            path += ".ply"

        try:
            points = self.point_cloud["points"]
            colors = self.point_cloud["colors"]

            point_cloud = trimesh.PointCloud(vertices=points)

            if colors is not None:
                colors_uint8 = (colors[:, :3] * 255).astype(np.uint8)
                point_cloud.colors = colors_uint8

            point_cloud.export(path)
            self.status_label.setText(f"点云已保存到: {path}")
        except Exception as exc:
            self.status_label.setText(f"保存失败: {exc}")

import os
import numpy as np
from PyQt5 import QtWidgets, QtCore, QtGui
import pyqtgraph.opengl as gl

from PyQt5.QtGui import QVector3D
import trimesh  # 用于点云采样


class PointCloudDialog(QtWidgets.QDialog):
    def __init__(self, parent=None, obj_data=None, groups=None):
        super().__init__(parent)
        self.setWindowTitle("点云采样")
        self.setGeometry(200, 200, 500, 400)
        self.obj_data = obj_data
        self.groups = groups
        self.point_cloud = None

        self.init_ui()

    def init_ui(self):
        layout = QtWidgets.QVBoxLayout(self)

        # 文件操作按钮
        file_btn_layout = QtWidgets.QHBoxLayout()

        self.import_btn = QtWidgets.QPushButton("加载 OBJ")
        self.import_btn.setIcon(self.style().standardIcon(QtWidgets.QStyle.SP_DialogOpenButton))
        self.import_btn.clicked.connect(self.import_obj)
        file_btn_layout.addWidget(self.import_btn)

        self.sample_btn = QtWidgets.QPushButton("采样点云")
        self.sample_btn.clicked.connect(self.sample_point_cloud)
        file_btn_layout.addWidget(self.sample_btn)

        layout.addLayout(file_btn_layout)

        # 点密度设置
        density_layout = QtWidgets.QHBoxLayout()
        density_layout.addWidget(QtWidgets.QLabel("点密度 (点/单位面积):"))
        self.density_spinbox = QtWidgets.QSpinBox()
        self.density_spinbox.setRange(1, 1000)
        self.density_spinbox.setValue(50)
        density_layout.addWidget(self.density_spinbox)
        layout.addLayout(density_layout)

        # 采样分辨率设置
        resolution_layout = QtWidgets.QHBoxLayout()
        resolution_layout.addWidget(QtWidgets.QLabel("采样分辨率:"))
        self.resolution_spinbox = QtWidgets.QSpinBox()
        self.resolution_spinbox.setRange(1, 1000)
        self.resolution_spinbox.setValue(100)
        resolution_layout.addWidget(self.resolution_spinbox)
        layout.addLayout(resolution_layout)

        # RGB选项
        self.rgb_checkbox = QtWidgets.QCheckBox("包含RGB颜色")
        self.rgb_checkbox.setChecked(True)
        layout.addWidget(self.rgb_checkbox)

        # 按钮布局
        button_layout = QtWidgets.QHBoxLayout()

        self.visualize_btn = QtWidgets.QPushButton("可视化")
        self.visualize_btn.clicked.connect(self.visualize_point_cloud)
        self.visualize_btn.setEnabled(False)
        button_layout.addWidget(self.visualize_btn)

        self.save_btn = QtWidgets.QPushButton("保存点云")
        self.save_btn.clicked.connect(self.save_point_cloud)
        self.save_btn.setEnabled(False)
        button_layout.addWidget(self.save_btn)

        layout.addLayout(button_layout)

        # 状态标签
        self.status_label = QtWidgets.QLabel("就绪")
        self.status_label.setStyleSheet("color: #AAAAAA; font-size: 10px;")
        layout.addWidget(self.status_label)

        # 点云视图
        self.pc_view = gl.GLViewWidget()
        self.pc_view.setMinimumSize(400, 300)
        self.pc_view.setBackgroundColor('#2D2D2D')
        layout.addWidget(self.pc_view)

    def import_obj(self):
        """导入新的OBJ文件"""
        path, _ = QtWidgets.QFileDialog.getOpenFileName(
            self, "打开 OBJ 文件", "", "OBJ 文件 (*.obj)")
        if not path:
            return

        try:
            # 解析OBJ文件
            vertices = []
            faces = []

            with open(path, 'r') as f:
                for line in f:
                    if line.startswith('v '):
                        parts = line.strip().split()
                        vertices.append([float(p) for p in parts[1:]])
                    elif line.startswith('f '):
                        parts = line.strip().split()
                        face = []
                        for part in parts[1:]:
                            v_idx = int(part.split('/')[0]) - 1
                            face.append(v_idx)
                        faces.append(face)

            if not vertices or not faces:
                self.status_label.setText("无效的 OBJ 文件 - 缺少顶点或面数据")
                return

            self.obj_data = {
                'vertices': np.array(vertices),
                'faces': faces,
                'path': path
            }

            # 重置组信息（可选）
            self.groups = None

            # 清空之前的点云
            self.point_cloud = None
            self.pc_view.clear()

            self.status_label.setText(f"已导入: {os.path.bename(path)}")
            self.visualize_btn.setEnabled(False)
            self.save_btn.setEnabled(False)

        except Exception as e:
            self.status_label.setText(f"导入失败: {str(e)}")

    def sample_point_cloud(self):
        if not self.obj_data or not self.obj_data['faces']:
            self.status_label.setText("错误: 没有可用的模型数据")
            return

        try:
            self.status_label.setText("正在采样点云...")
            QtWidgets.QApplication.processEvents()

            # 创建三角形网格
            vertices = self.obj_data['vertices']
            faces = self.obj_data['faces']

            # 创建trimesh对象
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)

            # 计算点密度
            point_density = self.density_spinbox.value()

            # 计算采样点数
            area = mesh.area
            num_points = int(area * point_density)

            # 确保点数合理
            if num_points < 100:
                num_points = 100
            elif num_points > 1000000:
                num_points = 1000000

            # 采样点云
            points, face_indices = trimesh.sample.sample_surface(mesh, num_points)

            # 如果需要颜色，从组件颜色中获取
            colors = None
            if self.rgb_checkbox.isChecked() and self.groups:
                # 创建面到组件的映射
                face_to_group = {}
                for group_idx, group in enumerate(self.groups):
                    for face in group['faces']:
                        # 找到原始面索引
                        for orig_face_idx, orig_face in enumerate(self.obj_data['faces']):
                            if set(orig_face) == set(group['v_indices'][v] for v in face):
                                face_to_group[orig_face_idx] = group_idx
                                break

                # 为每个点分配颜色
                colors = np.zeros((len(points), 4))
                for i, face_idx in enumerate(face_indices):
                    if face_idx in face_to_group:
                        group_idx = face_to_group[face_idx]
                        colors[i] = self.groups[group_idx]['color']
                    else:
                        colors[i] = (0.8, 0.8, 0.8, 1.0)  # 默认灰色
            elif self.rgb_checkbox.isChecked():
                # 没有组件信息，使用单一颜色
                colors = np.zeros((len(points), 4))
                colors[:, :3] = [0.8, 0.8, 0.8]  # 灰色
                colors[:, 3] = 1.0  # 完全不透明

            self.point_cloud = {
                'points': points,
                'colors': colors,
                'num_points': num_points
            }

            self.status_label.setText(f"采样完成! 点云包含 {num_points} 个点")
            self.visualize_btn.setEnabled(True)
            self.save_btn.setEnabled(True)

        except Exception as e:
            self.status_label.setText(f"采样失败: {str(e)}")

    def visualize_point_cloud(self):
        if not self.point_cloud:
            return

        self.pc_view.clear()

        points = self.point_cloud['points']
        colors = self.point_cloud['colors']

        # 创建点云对象
        point_cloud_item = gl.GLScatterPlotItem(
            pos=points,
            size=0.1,
            color=colors[:, :3] if colors is not None else (0.8, 0.8, 0.8, 1.0),
            pxMode=False
        )

        self.pc_view.addItem(point_cloud_item)

        # 添加坐标轴
        axis = gl.GLAxisItem()
        axis.setSize(10, 10, 10)
        self.pc_view.addItem(axis)

        # 自动调整视图
        if points.shape[0] > 0:
            min_coords = np.min(points, axis=0)
            max_coords = np.max(points, axis=0)
            size = np.max(max_coords - min_coords)
            center = (min_coords + max_coords) / 2
            distance = size * 2

            self.pc_view.setCameraPosition(
                pos=QVector3D(float(center[0]), float(center[1]), float(center[2] + distance)),
                distance=distance,
                elevation=30,
                azimuth=45
            )

    def save_point_cloud(self):
        if not self.point_cloud:
            return

        default_path = ""
        if self.obj_data and 'path' in self.obj_data:
            original_path = self.obj_data['path']
            dirname = os.path.dirname(original_path)
            basename = os.path.splitext(os.path.basename(original_path))[0]
            default_path = os.path.join(dirname, f"{basename}_pointcloud.ply")

        path, _ = QtWidgets.QFileDialog.getSaveFileName(
            self, "保存点云", default_path, "点云文件 (*.ply)")

        if not path:
            return

        # 确保文件扩展名
        if not path.lower().endswith('.ply'):
            path += '.ply'

        try:
            points = self.point_cloud['points']
            colors = self.point_cloud['colors']

            # 创建点云网格
            point_cloud = trimesh.PointCloud(vertices=points)

            # 添加颜色（如果有）
            if colors is not None:
                # 转换为0-255的整数
                colors_uint8 = (colors[:, :3] * 255).astype(np.uint8)
                point_cloud.colors = colors_uint8

            # 保存为PLY文件
            point_cloud.export(path)

            self.status_label.setText(f"点云已保存到: {path}")

        except Exception as e:
            self.status_label.setText(f"保存失败: {str(e)}")